import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms


# 修改 MyModule 类，实现卷积操作
class MyModule(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super().__init__()
        # 定义卷积层：nn.Conv2d(输入通道数, 输出通道数, 卷积核大小)
        self.conv_layer = nn.Conv2d(
            in_channels=in_channels,  # 输入图像的通道数（如3通道RGB）
            out_channels=out_channels,  # 输出特征图的通道数
            kernel_size=kernel_size,  # 卷积核大小（如3x3）
            stride=1  # 步长（默认1）
        )
        # （可选）如果需要使用自定义卷积核，可手动设置权重
        # 例如：用之前的自定义核初始化
        custom_kernel = torch.tensor([[1,2,1],[0,1,0],[2,1,0]], dtype=torch.float32)
        custom_kernel = custom_kernel.repeat(out_channels, in_channels, 1, 1)  # 适配多通道
        self.conv_layer.weight.data = custom_kernel

    def forward(self, x):
        # 前向传播：输入通过卷积层
        output = self.conv_layer(x)
        return output


class Mydata(Dataset):
    def __init__(self, root_dir, label_dir, transform=None):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir, self.label_dir)
        self.image_path = os.listdir(self.path)
        self.transform = transform

    def __getitem__(self, idx):
        image_name = self.image_path[idx]
        image_item_path = os.path.join(self.path, image_name)
        image = Image.open(image_item_path)
        label = self.label_dir
        if self.transform:
            image = self.transform(image)
        return image, label

    def __len__(self):
        return len(self.image_path)


# 数据预处理
dataset_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

# 加载数据集
train_root_dir = "train"
ants_label_dir = "ants"
ants_dataset = Mydata(train_root_dir, ants_label_dir, transform=dataset_transform)
test_loader = DataLoader(dataset=ants_dataset, batch_size=1, shuffle=True, num_workers=0)

# 实例化 MyModule：输入3通道（RGB），输出3通道，卷积核3x3
my_conv_module = MyModule(in_channels=3, out_channels=3, kernel_size=3)

# 可视化卷积结果
writer = SummaryWriter("conv_module_logs")
step = 0
for data in test_loader:
    ants_imgs, ants_labels = data  # 形状：[1, 3, 256, 256]

    # 使用 MyModule 执行卷积
    ants_imgs_conv = my_conv_module(ants_imgs)  # 输出形状：[1, 3, 254, 254]

    # 可视化
    writer.add_images("conv_result_by_module", ants_imgs_conv, step)
    step += 1

writer.close()